Skip to content

fix(api): add bulk_upsert method and unit tests for opensearch BaseOSDB class - #990

Draft
vkuznet wants to merge 14 commits into
DIRACGrid:mainfrom
vkuznet:fix/issue-401
Draft

fix(api): add bulk_upsert method and unit tests for opensearch BaseOSDB class#990
vkuznet wants to merge 14 commits into
DIRACGrid:mainfrom
vkuznet:fix/issue-401

Conversation

@vkuznet

@vkuznet vkuznet commented Jul 14, 2026

Copy link
Copy Markdown

This PR fixes issue #401 by introducing the bulk_upsert API to BaseOSDB class. It is also complement by full set of unit tests for BaseOSDB class which were missing. The unit test introduces mock client and different classes for different use cases.

Depends on #1007

Please note: it is my first PR and I'm happy to adjust it according to requirements/guidelines of DiracX community.

@aldbr aldbr linked an issue Jul 15, 2026 that may be closed by this pull request

@fstagni fstagni left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the original issue (#401) it was mentioned a specific use case. Can you already verify if this fix covers that? You can add it directly in this PR.

Comment thread diracx-db/src/diracx/db/os/utils.py
@vkuznet

vkuznet commented Aug 6, 2026

Copy link
Copy Markdown
Author

In the original issue (#401) it was mentioned a specific use case. Can you already verify if this fix covers that? You can add it directly in this PR.

@fstagni regarding issue #401. I checked the main repository code and what is described in this issue is not there the issue is based on commit outside of any branch, most likely in fork outside of the repository

# TODO: can we upsert to multiple documents?
for job_id, p_updates_ in param_updates.items():
if p_updates_:
await job_parameters_db.upsert(
int(job_id),
p_updates_,
)
). Therefore, before adding anything to this PR with usage of bulk insert I need to know if it is in fact required by code from main repository, see utils.py from main branch

@vkuznet vkuznet self-assigned this Aug 6, 2026
@aldbr
aldbr requested a review from fstagni August 6, 2026 14:36
@fstagni

fstagni commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

Hi, since #410 was created, the mentioned code has been moved and refactored, and can now be found in https://github.com/DIRACGrid/diracx/blob/main/diracx-logic/src/diracx/logic/jobs/status.py#L643. The issue anyway still holds, can you have a look?

@aldbr aldbr left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @vkuznet!
Just a few additional comments, sorry for the delay we take to review PRs...

response,
)

async def bulk_upsert(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

bulk_upsert would need to overridden within job_parameters_db because we are inserting a JobID and a timestamp:

def upsert(self, vo, doc_id, document):
document = {
"JobID": doc_id,
"timestamp": int(datetime.now(tz=UTC).timestamp() * 1000),
**document,
}
return super().upsert(vo, doc_id, document)

Here it would not work I think (and it looks like it's not spotted within the tests).

I actually wonder whether upsert is useful now that we have bulk_upsert.
I would suggest we just drop upsert and replace it everywhere with bulk_upsert, what do you think?

if new_application:
job_data["ApplicationStatus"] = new_application

await job_parameters_db.upsert(res["VO"], job_id, {"Status": new_status})

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

bulk_upsert could (should) be used here too I think

self.client,
actions,
raise_on_error=False,
raise_on_exception=False,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm just wondering what happens if there is a connection issue with the DB and no exception is raised.
I guess you would get 0 success, N errors but would get any information to know that there is an issue with the DB itself?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Comment on lines +683 to +686
if errors:
for error in errors:
logger.error("bulk insert error %s", error)
raise DocumentUpsertError("Failed to perform bulk insert operation")

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I assume this piece of code is generic because will be reused every time there is an error.
Wouldn't it make sense to raise the DocumentUpsertError from diracx-db itself? So that this part of the code only lives in db/os/utils and is automatically reused by the callers

Comment on lines +39 to +49
class DummyOSDB(BaseOSDB):
fields = {
"job_id": {"type": "long"},
"status": {"type": "keyword"},
"timestamp": {"type": "date"},
"vo": {"type": "keyword"},
}
index_prefix = "dummy"

def index_name(self, vo: str, doc_id: int) -> str:
return f"{self.index_prefix}-{vo}-{doc_id % 10}"

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Any reason for not reusing diracx-testing DummyOSDB? https://github.com/DIRACGrid/diracx/blob/d72016500fb6cfd73350aa7b2582c5dd2cdcf2d2/diracx-testing/src/diracx/testing/dummy_osdb.py

Like in

async def test_specified_parameters(prefilled_db: DummyOSDB):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think you can remove this conftest.py file and add your mock_osdb fixture within os_db.py, which is already including in pyproject.toml of the diracx-* packages. See in

"-pdiracx.testing", "-pdiracx.testing.osdb",

Though it looks like this is not documented...

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same comment here

@DIRACGridBot
DIRACGridBot marked this pull request as draft September 2, 2026 11:05
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Support upserts to multiple documents for opensearch

3 participants